# -*- coding: utf-8 -*-
"""
交易日历工具
"""

from datetime import date, datetime, timedelta
import logging
from typing import Set


class TradeCalendar:
    """交易日历类"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
        # 2024年节假日（可以根据实际情况更新）
        self.holidays_2024 = {
            # 元旦
            date(2024, 1, 1),
            
            # 春节
            date(2024, 2, 10), date(2024, 2, 11), date(2024, 2, 12),
            date(2024, 2, 13), date(2024, 2, 14), date(2024, 2, 15),
            date(2024, 2, 16), date(2024, 2, 17),
            
            # 清明节
            date(2024, 4, 4), date(2024, 4, 5), date(2024, 4, 6),
            
            # 劳动节
            date(2024, 5, 1), date(2024, 5, 2), date(2024, 5, 3),
            date(2024, 5, 4), date(2024, 5, 5),
            
            # 端午节
            date(2024, 6, 10),
            
            # 中秋节
            date(2024, 9, 15), date(2024, 9, 16), date(2024, 9, 17),
            
            # 国庆节
            date(2024, 10, 1), date(2024, 10, 2), date(2024, 10, 3),
            date(2024, 10, 4), date(2024, 10, 5), date(2024, 10, 6),
            date(2024, 10, 7),
        }
        
        # 2025年节假日（根据需要添加）
        self.holidays_2025 = {
            # 元旦
            date(2025, 1, 1),
            
            # 春节（示例，需要根据实际日历更新）
            date(2025, 1, 28), date(2025, 1, 29), date(2025, 1, 30),
            date(2025, 1, 31), date(2025, 2, 1), date(2025, 2, 2),
            date(2025, 2, 3),
        }
        
        # 合并所有节假日
        self.all_holidays = self.holidays_2024.union(self.holidays_2025)
    
    def is_trade_day(self, check_date: date = None) -> bool:
        """
        判断是否为交易日
        交易日定义：周一至周五，且不是节假日
        """
        if check_date is None:
            check_date = date.today()
        
        # 检查是否为周末
        if check_date.weekday() >= 5:  # 5=Saturday, 6=Sunday
            return False
        
        # 检查是否为节假日
        if check_date in self.all_holidays:
            return False
        
        return True
    
    def get_next_trade_day(self, start_date: date = None) -> date:
        """获取下一个交易日"""
        if start_date is None:
            start_date = date.today()
        
        next_day = start_date + timedelta(days=1)
        while not self.is_trade_day(next_day):
            next_day += timedelta(days=1)
        
        return next_day
    
    def get_previous_trade_day(self, start_date: date = None) -> date:
        """获取前一个交易日"""
        if start_date is None:
            start_date = date.today()
        
        previous_day = start_date - timedelta(days=1)
        while not self.is_trade_day(previous_day):
            previous_day -= timedelta(days=1)
        
        return previous_day
    
    def get_trade_days_in_range(self, start_date: date, end_date: date) -> list:
        """获取指定日期范围内的所有交易日"""
        trade_days = []
        current_date = start_date
        
        while current_date <= end_date:
            if self.is_trade_day(current_date):
                trade_days.append(current_date)
            current_date += timedelta(days=1)
        
        return trade_days
    
    def count_trade_days(self, start_date: date, end_date: date) -> int:
        """计算指定日期范围内的交易日数量"""
        return len(self.get_trade_days_in_range(start_date, end_date))
    
    def is_trade_time(self, check_time: datetime = None) -> bool:
        """
        判断是否为交易时间
        交易时间：交易日的 9:30-15:00
        """
        if check_time is None:
            check_time = datetime.now()
        
        # 首先检查是否为交易日
        if not self.is_trade_day(check_time.date()):
            return False
        
        # 检查时间是否在交易时间内
        trade_start = check_time.replace(hour=9, minute=30, second=0, microsecond=0)
        trade_end = check_time.replace(hour=15, minute=0, second=0, microsecond=0)
        
        return trade_start <= check_time <= trade_end
    
    def get_next_buy_time(self) -> datetime:
        """获取下一个买入时间（下一个交易日的9:30）"""
        next_trade_day = self.get_next_trade_day()
        return datetime.combine(next_trade_day, datetime.min.time().replace(hour=9, minute=30))
    
    def get_next_sell_time(self) -> datetime:
        """获取下一个收盘时间（下一个交易日的15:00）"""
        next_trade_day = self.get_next_trade_day()
        return datetime.combine(next_trade_day, datetime.min.time().replace(hour=15, minute=0))
    
    def add_holiday(self, holiday_date: date):
        """添加节假日"""
        self.all_holidays.add(holiday_date)
        self.logger.info(f"已添加节假日: {holiday_date}")
    
    def remove_holiday(self, holiday_date: date):
        """移除节假日"""
        if holiday_date in self.all_holidays:
            self.all_holidays.remove(holiday_date)
            self.logger.info(f"已移除节假日: {holiday_date}")
    
    def get_holiday_list(self) -> list:
        """获取所有节假日列表"""
        return sorted(list(self.all_holidays))
